# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#     http://www.apache.org/licenses/LICENSE-2.0
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.

"""
CustNameService defines a CustName text matching service
"""
import json
import logging
import traceback
from bert4keras.snippets import sequence_padding

from model_service import ModelService

class CustNameService(ModelService):
    """
    CustNameService defines a fundamental service for CustName text matching.
    """

    def preprocess(self, request):
        """
        Decode all input images into ndarray.

        Note: This implementation doesn't properly handle error cases in batch mode,
        If one of the input images is corrupted, all requests in the batch will fail.

        :param request:
        :return:
        """
        param_name = self.signature['inputs'][0]['data_name']

        for idx, data in enumerate(request):
            text = data.get(param_name)
            if text is None:
                text = data.get("body")

            if text is None:
                text = data.get("data")

            if text is None or len(text) == 0:
                self.error = "Empty text input"
                return None

            text = text.decode('utf-8') # bytearray转为string
            text = eval(text) # 字符串转成列表
            #logging.info("Begin text matching %s" % text)

            try:
                # 格式化文本
                batch_token_ids, batch_segment_ids = [], []
                for (text1, text2) in text:
                    token_ids, segment_ids = self.tokenizer.encode(
                        text1, text2, maxlen=self.maxlen
                    )
                    batch_token_ids.append(token_ids)
                    batch_segment_ids.append(segment_ids)
                # 补齐
                batch_token_ids = sequence_padding(batch_token_ids)
                batch_segment_ids = sequence_padding(batch_segment_ids)
                # 预测
                y_pred = self.model.predict([batch_token_ids, batch_segment_ids]).argmax(axis=1)
            except Exception as e:
                logging.error('Text matching error:%s, errmsg:%s' % (text, traceback.format_exc()))
                self.error = "data must be [(text1, text2)] fromat!"
                return None

            # y_pred = [int(i) for i in y_pred]

        return y_pred

    def postprocess(self, data):
        if self.error is not None:
            return [self.error] * self._batch_size

        return [str(data)]

_service = CustNameService()


def handle(data, context):
    if not _service.initialized:
        _service.initialize(context)

    if data is None:
        return None

    return _service.handle(data, context)
